import numpy as np
import matplotlib.pyplot as plt

# 设置参数
n = 64  # 网格点数（内部点）
k_values = np.linspace(1, n, 1000)  # 波数范围1到63

# 计算特征值和迭代次数
def compute_iterations(omega, n, k_values):
    theta = k_values * np.pi / (2 * n)
    sin_sq = np.sin(theta) ** 2
    lambdas = 1 - 2 * omega * sin_sq
    iterations = []
    for l in np.abs(lambdas):
        if l == 0:
            iterations.append(1)
        elif l >= 1:
            iterations.append(100)
        else:
            it = np.log(0.01) / np.log(l)
            if (it <= 100) : iterations.append(np.ceil(it))
            else : iterations.append(100)
    return np.array(iterations)

# 计算两种omega的情况
omega1_iters = compute_iterations(1.0, n, k_values)
omega23_iters = compute_iterations(2/3, n, k_values)

# 绘制图形
plt.figure(figsize=(12, 5))

# Omega=1的情况
plt.subplot(1, 2, 1)
plt.plot(k_values, omega1_iters, 'b-')
plt.title('Weighted Jacobi (ω=1)')
plt.xlabel('Wavenumber k')
plt.ylabel('Iterations Required')
plt.grid(True)
plt.ylim(0, 100)

# Omega=2/3的情况
plt.subplot(1, 2, 2)
plt.plot(k_values, omega23_iters, 'r-')
plt.title('Weighted Jacobi (ω=2/3)')
plt.xlabel('Wavenumber k')
plt.grid(True)
plt.ylim(0, 100) 

plt.tight_layout()
plt.savefig("./E9_21.jpg")